
"""
可视化工具

将数据转换为表格和图表在Chainlit中展示。
"""

import chainlit as cl
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import logging
from typing import Dict, Any, List, Optional, Union, Literal
from datetime import datetime


# 显示模式类型
ElementDisplay = Literal["inline", "side", "page"]


class Visualizer:
    """数据可视化工具"""
    
    def __init__(self):
        """初始化可视化工具"""
        self.logger = logging.getLogger("ui.visualizer")
    
    async def display_table(
        self,
        df: pd.DataFrame,
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示表格
        
        Args:
            df: 要显示的DataFrame
            title: 表格标题
            display_mode: 显示模式（inline或side）
        """
        try:
            if df.empty:
                await cl.Message(content=f"{title}\n\n暂无数据").send()
                return
            
            # 将DataFrame转换为markdown格式
            markdown_table = df.to_markdown(index=False)
            
            # 创建文本元素
            elements = [
                cl.Text(
                    name=title or "数据表格",
                    content=markdown_table,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {title}" if title else "### 数据表格",
                elements=elements
            ).send()
            
            self.logger.info(f"显示表格成功: {title}, 行数={len(df)}")
            
        except Exception as e:
            self.logger.error(f"显示表格失败: {e}")
            await cl.Message(content=f"显示表格失败: {str(e)}").send()
    
    async def display_dataframe(
        self,
        data: Union[pd.DataFrame, Dict[str, Any], List[Dict[str, Any]]],
        title: str = "",
        max_rows: int = 50
    ):
        """
        显示DataFrame（支持多种数据格式）
        
        Args:
            data: 数据（DataFrame、字典或字典列表）
            title: 标题
            max_rows: 最大显示行数
        """
        try:
            # 转换为DataFrame
            if isinstance(data, pd.DataFrame):
                df = data
            elif isinstance(data, dict):
                df = pd.DataFrame([data])
            elif isinstance(data, list):
                df = pd.DataFrame(data)
            else:
                await cl.Message(content=f"{title}\n\n不支持的数据格式").send()
                return
            
            # 限制行数
            if len(df) > max_rows:
                df = df.head(max_rows)
                title = f"{title} (仅显示前{max_rows}行)"
            
            await self.display_table(df, title)
            
        except Exception as e:
            self.logger.error(f"显示DataFrame失败: {e}")
            await cl.Message(content=f"显示数据失败: {str(e)}").send()
    
    async def display_chart(
        self,
        data: Dict[str, Any],
        chart_type: str = "bar",
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示图表
        
        Args:
            data: 图表数据
            chart_type: 图表类型（bar, line, scatter, pie等）
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            fig = None
            
            if chart_type == "bar":
                fig = go.Figure(data=[
                    go.Bar(
                        x=data.get("x", []),
                        y=data.get("y", []),
                        name=data.get("name", ""),
                        marker_color=data.get("color", "#1f77b4")
                    )
                ])
            
            elif chart_type == "line":
                fig = go.Figure(data=[
                    go.Scatter(
                        x=data.get("x", []),
                        y=data.get("y", []),
                        mode='lines+markers',
                        name=data.get("name", ""),
                        line=dict(color=data.get("color", "#1f77b4"))
                    )
                ])
            
            elif chart_type == "scatter":
                fig = go.Figure(data=[
                    go.Scatter(
                        x=data.get("x", []),
                        y=data.get("y", []),
                        mode='markers',
                        name=data.get("name", ""),
                        marker=dict(
                            size=data.get("size", 10),
                            color=data.get("color", "#1f77b4")
                        )
                    )
                ])
            
            elif chart_type == "pie":
                fig = go.Figure(data=[
                    go.Pie(
                        labels=data.get("labels", []),
                        values=data.get("values", []),
                        hole=data.get("hole", 0)
                    )
                ])
            
            else:
                await cl.Message(content=f"不支持的图表类型: {chart_type}").send()
                return
            
            # 更新布局
            fig.update_layout(
                title=title or data.get("title", ""),
                xaxis_title=data.get("x_label", ""),
                yaxis_title=data.get("y_label", ""),
                height=data.get("height", 400),
                showlegend=data.get("show_legend", True)
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=title or "图表",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {title}" if title else "",
                elements=elements
            ).send()
            
            self.logger.info(f"显示图表成功: {title}, 类型={chart_type}")
            
        except Exception as e:
            self.logger.error(f"显示图表失败: {e}")
            await cl.Message(content=f"显示图表失败: {str(e)}").send()
    
    async def display_multi_series_chart(
        self,
        data: List[Dict[str, Any]],
        chart_type: str = "line",
        title: str = "",
        x_label: str = "",
        y_label: str = ""
    ):
        """
        显示多系列图表
        
        Args:
            data: 多个系列的数据列表
            chart_type: 图表类型
            title: 图表标题
            x_label: X轴标签
            y_label: Y轴标签
        """
        try:
            fig = go.Figure()
            
            for series in data:
                if chart_type == "line":
                    fig.add_trace(go.Scatter(
                        x=series.get("x", []),
                        y=series.get("y", []),
                        mode='lines+markers',
                        name=series.get("name", ""),
                        line=dict(color=series.get("color"))
                    ))
                elif chart_type == "bar":
                    fig.add_trace(go.Bar(
                        x=series.get("x", []),
                        y=series.get("y", []),
                        name=series.get("name", ""),
                        marker_color=series.get("color")
                    ))
            
            fig.update_layout(
                title=title,
                xaxis_title=x_label,
                yaxis_title=y_label,
                height=400,
                showlegend=True
            )
            
            elements = [cl.Plotly(name=title, figure=fig, display="inline")]
            await cl.Message(content=f"### {title}", elements=elements).send()
            
            self.logger.info(f"显示多系列图表成功: {title}")
            
        except Exception as e:
            self.logger.error(f"显示多系列图表失败: {e}")
            await cl.Message(content=f"显示图表失败: {str(e)}").send()
    
    async def display_candlestick_chart(
        self,
        stock_code: str,
        data: pd.DataFrame,
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示股票蜡烛图（K线图）
        
        Args:
            stock_code: 股票代码
            data: 包含日期、开盘、最高、最低、收盘、成交量的DataFrame
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if data.empty:
                await cl.Message(content=f"{title}\n\n暂无K线数据").send()
                return
            
            # 创建蜡烛图
            fig = go.Figure(data=[
                go.Candlestick(
                    x=data['日期'],
                    open=data['开盘'],
                    high=data['最高'],
                    low=data['最低'],
                    close=data['收盘'],
                    name='K线'
                )
            ])
            
            # 添加成交量柱状图（子图）
            fig.add_trace(go.Bar(
                x=data['日期'],
                y=data['成交量'],
                name='成交量',
                marker_color='rgba(100, 100, 100, 0.5)',
                yaxis='y2'
            ))
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} K线图",
                xaxis_title="日期",
                yaxis_title="价格",
                height=600,
                showlegend=True,
                xaxis_rangeslider_visible=False,
                yaxis2=dict(
                    title="成交量",
                    overlaying='y',
                    side='right',
                    showgrid=False
                )
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} K线图",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} K线图分析" if title else f"### {stock_code} K线图",
                elements=elements
            ).send()
            
            self.logger.info(f"显示蜡烛图成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示蜡烛图失败: {e}")
            await cl.Message(content=f"显示K线图失败: {str(e)}").send()
    
    async def display_technical_indicators_chart(
        self,
        stock_code: str,
        data: pd.DataFrame,
        indicators: Dict[str, Any],
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示技术指标图表（MACD、RSI、布林带等）
        
        Args:
            stock_code: 股票代码
            data: 价格数据
            indicators: 技术指标数据
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if data.empty:
                await cl.Message(content=f"{title}\n\n暂无技术指标数据").send()
                return
            
            # 创建子图
            from plotly.subplots import make_subplots
            fig = make_subplots(
                rows=3, cols=1,
                shared_xaxes=True,
                vertical_spacing=0.05,
                subplot_titles=('价格走势', 'MACD指标', 'RSI指标'),
                row_heights=[0.5, 0.25, 0.25]
            )
            
            # 价格走势图
            fig.add_trace(
                go.Scatter(
                    x=data['日期'],
                    y=data['收盘'],
                    mode='lines',
                    name='收盘价',
                    line=dict(color='#1f77b4')
                ),
                row=1, col=1
            )
            
            # 添加移动平均线
            if 'ma5' in indicators and 'ma20' in indicators:
                fig.add_trace(
                    go.Scatter(
                        x=data['日期'],
                        y=indicators['ma5'],
                        mode='lines',
                        name='MA5',
                        line=dict(color='orange', width=1)
                    ),
                    row=1, col=1
                )
                fig.add_trace(
                    go.Scatter(
                        x=data['日期'],
                        y=indicators['ma20'],
                        mode='lines',
                        name='MA20',
                        line=dict(color='red', width=1)
                    ),
                    row=1, col=1
                )
            
            # MACD指标
            if 'macd' in indicators and 'macd_signal' in indicators:
                fig.add_trace(
                    go.Scatter(
                        x=data['日期'],
                        y=indicators['macd'],
                        mode='lines',
                        name='MACD',
                        line=dict(color='blue')
                    ),
                    row=2, col=1
                )
                fig.add_trace(
                    go.Scatter(
                        x=data['日期'],
                        y=indicators['macd_signal'],
                        mode='lines',
                        name='Signal',
                        line=dict(color='red')
                    ),
                    row=2, col=1
                )
                
                # MACD柱状图
                if 'macd_hist' in indicators:
                    colors = ['green' if x >= 0 else 'red' for x in indicators['macd_hist']]
                    fig.add_trace(
                        go.Bar(
                            x=data['日期'],
                            y=indicators['macd_hist'],
                            name='MACD Hist',
                            marker_color=colors
                        ),
                        row=2, col=1
                    )
            
            # RSI指标
            if 'rsi' in indicators:
                fig.add_trace(
                    go.Scatter(
                        x=data['日期'],
                        y=indicators['rsi'],
                        mode='lines',
                        name='RSI',
                        line=dict(color='purple')
                    ),
                    row=3, col=1
                )
                
                # 添加RSI参考线
                fig.add_hline(y=70, line_dash="dash", line_color="red", row="3", col="1")
                fig.add_hline(y=30, line_dash="dash", line_color="green", row="3", col="1")
                fig.add_hline(y=50, line_dash="dot", line_color="gray", row="3", col="1")
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} 技术指标分析",
                height=800,
                showlegend=True
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} 技术指标",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} 技术指标分析",
                elements=elements
            ).send()
            
            self.logger.info(f"显示技术指标图表成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示技术指标图表失败: {e}")
            await cl.Message(content=f"显示技术指标图表失败: {str(e)}").send()
    
    async def display_capital_flow_chart(
        self,
        stock_code: str,
        flow_data: Dict[str, Any],
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示资金流向堆叠图
        
        Args:
            stock_code: 股票代码
            flow_data: 资金流向数据
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if not flow_data:
                await cl.Message(content=f"{title}\n\n暂无资金流向数据").send()
                return
            
            # 创建资金流向堆叠图
            categories = ['主力资金', '散户资金', '其他资金']
            inflow_data = [
                flow_data.get('main_inflow', 0),
                flow_data.get('retail_inflow', 0),
                flow_data.get('other_inflow', 0)
            ]
            outflow_data = [
                flow_data.get('main_outflow', 0),
                flow_data.get('retail_outflow', 0),
                flow_data.get('other_outflow', 0)
            ]
            
            fig = go.Figure()
            
            # 流入资金
            fig.add_trace(go.Bar(
                name='资金流入',
                x=categories,
                y=inflow_data,
                marker_color='green'
            ))
            
            # 流出资金
            fig.add_trace(go.Bar(
                name='资金流出',
                x=categories,
                y=outflow_data,
                marker_color='red'
            ))
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} 资金流向分析",
                xaxis_title="资金类型",
                yaxis_title="金额(元)",
                barmode='stack',
                height=400
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} 资金流向",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} 资金流向分析",
                elements=elements
            ).send()
            
            self.logger.info(f"显示资金流向图表成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示资金流向图表失败: {e}")
            await cl.Message(content=f"显示资金流向图表失败: {str(e)}").send()
    
    async def display_financial_ratios_radar(
        self,
        stock_code: str,
        ratios: Dict[str, float],
        industry_averages: Optional[Dict[str, float]] = None,
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示财务比率雷达图
        
        Args:
            stock_code: 股票代码
            ratios: 财务比率数据
            industry_averages: 行业平均数据
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if not ratios:
                await cl.Message(content=f"{title}\n\n暂无财务比率数据").send()
                return
            
            # 定义雷达图维度
            categories = ['盈利能力', '偿债能力', '运营效率', '成长性', '现金流']
            
            # 计算各维度得分（简化处理）
            scores = []
            for category in categories:
                if category == '盈利能力':
                    score = ratios.get('roe', 0) / 20 * 100 if ratios.get('roe') else 50
                elif category == '偿债能力':
                    score = max(0, 100 - (ratios.get('debt_ratio', 0) or 0))
                elif category == '运营效率':
                    score = ratios.get('asset_turnover', 0) * 100 if ratios.get('asset_turnover') else 50
                elif category == '成长性':
                    score = min(100, ratios.get('revenue_growth', 0) * 10) if ratios.get('revenue_growth') else 50
                elif category == '现金流':
                    score = min(100, ratios.get('cash_flow_ratio', 0) * 20) if ratios.get('cash_flow_ratio') else 50
                else:
                    score = 50
                scores.append(score)
            
            # 创建雷达图
            fig = go.Figure()
            
            fig.add_trace(go.Scatterpolar(
                r=scores,
                theta=categories,
                fill='toself',
                name=stock_code
            ))
            
            # 添加行业平均线（如果有）
            if industry_averages:
                industry_scores = []
                for category in categories:
                    if category == '盈利能力':
                        score = industry_averages.get('roe', 0) / 20 * 100 if industry_averages.get('roe') else 50
                    elif category == '偿债能力':
                        score = max(0, 100 - (industry_averages.get('debt_ratio', 0) or 0))
                    elif category == '运营效率':
                        score = industry_averages.get('asset_turnover', 0) * 100 if industry_averages.get('asset_turnover') else 50
                    elif category == '成长性':
                        score = min(100, industry_averages.get('revenue_growth', 0) * 10) if industry_averages.get('revenue_growth') else 50
                    elif category == '现金流':
                        score = min(100, industry_averages.get('cash_flow_ratio', 0) * 20) if industry_averages.get('cash_flow_ratio') else 50
                    else:
                        score = 50
                    industry_scores.append(score)
                
                fig.add_trace(go.Scatterpolar(
                    r=industry_scores,
                    theta=categories,
                    fill='toself',
                    name='行业平均'
                ))
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} 财务比率雷达图",
                polar=dict(
                    radialaxis=dict(
                        visible=True,
                        range=[0, 100]
                    )
                ),
                showlegend=True,
                height=500
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} 财务比率雷达图",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} 财务比率分析",
                elements=elements
            ).send()
            
            self.logger.info(f"显示财务比率雷达图成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示财务比率雷达图失败: {e}")
            await cl.Message(content=f"显示财务比率雷达图失败: {str(e)}").send()
    
    async def display_price_volume_chart(
        self,
        stock_code: str,
        data: pd.DataFrame,
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示价格成交量组合图
        
        Args:
            stock_code: 股票代码
            data: 包含日期、收盘价、成交量的DataFrame
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if data.empty:
                await cl.Message(content=f"{title}\n\n暂无价格成交量数据").send()
                return
            
            # 创建子图
            from plotly.subplots import make_subplots
            fig = make_subplots(
                rows=2, cols=1,
                shared_xaxes=True,
                vertical_spacing=0.05,
                subplot_titles=('价格走势', '成交量'),
                row_heights=[0.7, 0.3]
            )
            
            # 价格走势图
            fig.add_trace(
                go.Scatter(
                    x=data['日期'],
                    y=data['收盘'],
                    mode='lines',
                    name='收盘价',
                    line=dict(color='#1f77b4')
                ),
                row=1, col=1
            )
            
            # 成交量柱状图
            colors = ['green' if data['收盘'].iloc[i] >= data['收盘'].iloc[i-1] else 'red'
                     for i in range(len(data))]
            fig.add_trace(
                go.Bar(
                    x=data['日期'],
                    y=data['成交量'],
                    name='成交量',
                    marker_color=colors
                ),
                row=2, col=1
            )
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} 价格成交量分析",
                height=600,
                showlegend=True
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} 价格成交量图",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} 价格成交量分析",
                elements=elements
            ).send()
            
            self.logger.info(f"显示价格成交量图表成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示价格成交量图表失败: {e}")
            await cl.Message(content=f"显示价格成交量图表失败: {str(e)}").send()
    
    async def display_bollinger_bands_chart(
        self,
        stock_code: str,
        data: pd.DataFrame,
        bollinger_data: Dict[str, Any],
        title: str = "",
        display_mode: ElementDisplay = "inline"
    ):
        """
        显示布林带图表
        
        Args:
            stock_code: 股票代码
            data: 价格数据
            bollinger_data: 布林带数据
            title: 图表标题
            display_mode: 显示模式
        """
        try:
            if data.empty:
                await cl.Message(content=f"{title}\n\n暂无布林带数据").send()
                return
            
            # 创建布林带图表
            fig = go.Figure()
            
            # 添加价格线
            fig.add_trace(go.Scatter(
                x=data['日期'],
                y=data['收盘'],
                mode='lines',
                name='收盘价',
                line=dict(color='#1f77b4')
            ))
            
            # 添加布林带上轨
            if 'upper_band' in bollinger_data:
                fig.add_trace(go.Scatter(
                    x=data['日期'],
                    y=bollinger_data['upper_band'],
                    mode='lines',
                    name='布林带上轨',
                    line=dict(color='red', dash='dash')
                ))
            
            # 添加布林带中轨
            if 'middle_band' in bollinger_data:
                fig.add_trace(go.Scatter(
                    x=data['日期'],
                    y=bollinger_data['middle_band'],
                    mode='lines',
                    name='布林带中轨',
                    line=dict(color='orange')
                ))
            
            # 添加布林带下轨
            if 'lower_band' in bollinger_data:
                fig.add_trace(go.Scatter(
                    x=data['日期'],
                    y=bollinger_data['lower_band'],
                    mode='lines',
                    name='布林带下轨',
                    line=dict(color='green', dash='dash')
                ))
            
            # 更新布局
            fig.update_layout(
                title=f"{stock_code} {title}" if title else f"{stock_code} 布林带分析",
                xaxis_title="日期",
                yaxis_title="价格",
                height=500,
                showlegend=True
            )
            
            # 创建Plotly元素
            elements = [
                cl.Plotly(
                    name=f"{stock_code} 布林带图",
                    figure=fig,
                    display=display_mode
                )
            ]
            
            # 发送消息
            await cl.Message(
                content=f"### {stock_code} 布林带分析",
                elements=elements
            ).send()
            
            self.logger.info(f"显示布林带图表成功: {stock_code}")
            
        except Exception as e:
            self.logger.error(f"显示布林带图表失败: {e}")
            await cl.Message(content=f"显示布林带图表失败: {str(e)}").send()


# 创建全局可视化器实例
visualizer = Visualizer()